mod arg_parser;
mod data;
mod model;
mod train;

use burn::{
    backend::{Autodiff, Wgpu},
    data::dataset::{Dataset, vision::MnistDataset},
    optim::AdamConfig,
};
use clap::Parser;

use crate::arg_parser::{Config, Operator};
use crate::{
    model::ModelConfig,
    train::{TrainingConfig, infer, train},
};

fn main() {
    type Backend = Autodiff<Wgpu<f32, i32>>;

    let device = Default::default();
    let artifact_dir = "results";
    let config = Config::parse();
    match config.op {
        Operator::Train {
            num_epochs,
            batch_size,
            learning_rate,
            seed,
        } => {
            let config = ModelConfig::new(1, 8, 16, 3, 8, 512, 10);
            let training_config = TrainingConfig::new(config, AdamConfig::new());
            let training_config = num_epochs.map_or(training_config.clone(), |item| {
                training_config.with_num_epochs(item)
            });
            let training_config = batch_size.map_or(training_config.clone(), |item| {
                training_config.with_batch_size(item)
            });
            let training_config = learning_rate.map_or(training_config.clone(), |item| {
                training_config.with_learning_rate(item)
            });
            let training_config = seed.map_or(training_config.clone(), |item| {
                training_config.with_seed(item)
            });
            train::<Backend>(artifact_dir, training_config, device);
        }
        Operator::Infer => {
            infer::<Backend>(artifact_dir, device, MnistDataset::test().get(42).unwrap());
        }
    }
}
